#!/usr/bin/env python3
import asyncio
import json
import logging
import websockets
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 定义消息类型常量
PATIENT_REPORT = 'PATIENT_REPORT'
PATIENT_REPORT_RESPONSE = 'PATIENT_REPORT_RESPONSE'
AUTH = 'AUTH'
HEARTBEAT = 'HEARTBEAT'
USER_REGISTER = 'USER_REGISTER'

class SimpleWebSocketServer:
    def __init__(self):
        self.clients = set()
        
    async def register(self, websocket):
        self.clients.add(websocket)
        logger.info(f"新客户端连接，当前连接数: {len(self.clients)}")
        
    async def unregister(self, websocket):
        self.clients.discard(websocket)
        logger.info(f"客户端断开连接，当前连接数: {len(self.clients)}")
        
    async def broadcast(self, message):
        if self.clients:
            await asyncio.gather(
                *[client.send(message) for client in self.clients],
                return_exceptions=True
            )
            
    async def handler(self, websocket, path):
        await self.register(websocket)
        try:
            async for message in websocket:
                try:
                    logger.info(f"收到原始消息: {message}")
                    data = json.loads(message)
                    logger.info(f"解析后的消息: {data}")
                    
                    message_type = data.get('type')
                    if not message_type:
                        logger.error("消息缺少type字段")
                        continue
                    
                    # 处理认证消息
                    if message_type == AUTH:
                        logger.info("收到认证消息")
                        try:
                            response = {
                                'type': 'AUTH_RESPONSE',
                                'data': {
                                    'status': 'success',
                                    'timestamp': datetime.now().isoformat()
                                }
                            }
                            await websocket.send(json.dumps(response))
                            logger.info("认证响应发送成功")
                        except Exception as e:
                            logger.error(f"发送认证响应失败: {e}")
                    
                    # 处理心跳消息
                    elif message_type == HEARTBEAT:
                        logger.info("收到心跳消息")
                        try:
                            response = {
                                'type': 'HEARTBEAT_RESPONSE',
                                'data': {
                                    'timestamp': datetime.now().isoformat()
                                }
                            }
                            await websocket.send(json.dumps(response))
                            logger.info("心跳响应发送成功")
                        except Exception as e:
                            logger.error(f"发送心跳响应失败: {e}")
                    
                    # 处理患者上报
                    elif message_type == PATIENT_REPORT:
                        logger.info("收到患者上报消息")
                        try:
                            response = {
                                'type': PATIENT_REPORT_RESPONSE,
                                'data': {
                                    'status': 'success',
                                    'timestamp': datetime.now().isoformat()
                                }
                            }
                            await self.broadcast(json.dumps(response))
                            logger.info("患者上报响应广播成功")
                        except Exception as e:
                            logger.error(f"广播患者上报响应失败: {e}")
                    
                    # 处理用户注册
                    elif message_type == USER_REGISTER:
                        logger.info("收到用户注册消息")
                        try:
                            response = {
                                'type': 'USER_REGISTER_RESPONSE',
                                'data': {
                                    'status': 'success',
                                    'timestamp': datetime.now().isoformat()
                                }
                            }
                            await self.broadcast(json.dumps(response))
                            logger.info("用户注册响应广播成功")
                        except Exception as e:
                            logger.error(f"广播用户注册响应失败: {e}")
                    else:
                        logger.warning(f"未知消息类型: {message_type}")
                        
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析错误: {e}, 消息: {message}")
                except Exception as e:
                    logger.error(f"处理消息时发生错误: {e}")
                    
        except websockets.exceptions.ConnectionClosed:
            logger.info("连接正常关闭")
        except Exception as e:
            logger.error(f"WebSocket处理过程中发生错误: {e}")
        finally:
            await self.unregister(websocket)

async def main():
    server = SimpleWebSocketServer()
    host = "0.0.0.0"
    port = 8080
    logger.info(f"启动WebSocket服务器: ws://{host}:{port}")
    
    # 使用正确的handler调用方式
    async with websockets.serve(server.handler, host, port):
        logger.info("WebSocket服务器启动成功")
        await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("服务器停止")
    except Exception as e:
        logger.error(f"服务器运行出错: {e}") 